package com.bfl.mybatis.generator;

import org.joda.time.LocalDate;
import org.joda.time.LocalDateTime;
import org.joda.time.LocalTime;
import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
import org.mybatis.generator.internal.types.JavaTypeResolverDefaultImpl;

import java.math.BigDecimal;
import java.sql.Types;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

import static org.mybatis.generator.internal.util.StringUtility.isTrue;

/**
 * https://github.com/agwlvssainokuni/mybatis-generator-custom
 */
public class CustomJavaTypeResolver extends JavaTypeResolverDefaultImpl {

    private static final String PROP_TINYINT1_IS_BIT = "tinyInt1isBit";
    private static final String PROP_USE_JODA_TIME = "useJodaTime";
    private static final String PROP_JAVA_TYPE_BY_COLUMN_NAME_PREFIX = "javaTypeByColumnName.";
    private static final String PROP_JAVA_TYPE_BY_JDBC_TYPE_NAME_PREFIX = "javaTypeByJdbcTypeName.";
    private boolean tinyInt1isBit;
    private boolean useJodaTime;
    private Map<String, FullyQualifiedJavaType> javaTypeByColumnName = new HashMap<>();
    private Map<String, FullyQualifiedJavaType> javaTypeByJdbcTypeName = new HashMap<>();

    public CustomJavaTypeResolver() {
        super();
        tinyInt1isBit = Boolean.TRUE;
        useJodaTime = Boolean.FALSE;
    }

    @Override
    public void addConfigurationProperties(Properties properties) {
        super.addConfigurationProperties(properties);
        tinyInt1isBit = isTrue(properties.getProperty(PROP_TINYINT1_IS_BIT, Boolean.TRUE.toString()));
        useJodaTime = isTrue(properties.getProperty(PROP_USE_JODA_TIME, Boolean.FALSE.toString()));

        if (useJodaTime) {
            typeMap.put(Types.DATE, new JdbcTypeInformation("DATE",
                    new FullyQualifiedJavaType(LocalDate.class.getName())));
            typeMap.put(Types.TIME, new JdbcTypeInformation("TIME",
                    new FullyQualifiedJavaType(LocalTime.class.getName())));
            typeMap.put(Types.TIMESTAMP, new JdbcTypeInformation("TIMESTAMP",
                    new FullyQualifiedJavaType(LocalDateTime.class.getName())));
        }

        for (String propName : properties.stringPropertyNames()) {
            if (propName.startsWith(PROP_JAVA_TYPE_BY_COLUMN_NAME_PREFIX)) {
                String columnName = propName.substring(PROP_JAVA_TYPE_BY_COLUMN_NAME_PREFIX.length());
                String javaType = properties.getProperty(propName);
                javaTypeByColumnName.put(columnName.toUpperCase(), new FullyQualifiedJavaType(javaType));
            }

            if (propName.startsWith(PROP_JAVA_TYPE_BY_JDBC_TYPE_NAME_PREFIX)) {
                String jdbcTypeName = propName.substring(PROP_JAVA_TYPE_BY_JDBC_TYPE_NAME_PREFIX.length());
                String javaType = properties.getProperty(propName);
                javaTypeByJdbcTypeName.put(jdbcTypeName.toUpperCase(), new FullyQualifiedJavaType(javaType));
            }
        }

    }

    @Override
    public FullyQualifiedJavaType calculateJavaType(IntrospectedColumn introspectedColumn) {

        FullyQualifiedJavaType result;

        // 通过 javaTypeByJdbcType.{JDBC_TYPE_NAME} 自定义类型映射
        String jdbcTypeName = introspectedColumn.getJdbcTypeName().toUpperCase();
        if (javaTypeByJdbcTypeName.containsKey(jdbcTypeName)) {
            return javaTypeByJdbcTypeName.get(jdbcTypeName);
        }

        // 通过 javaTypeByColumnName.{COLUMN_NAME} 自定义类型映射
        String columnName = introspectedColumn.getActualColumnName().toUpperCase();
        if (javaTypeByColumnName.containsKey(columnName)) {
            return javaTypeByColumnName.get(columnName);
        }

        // TINYINT(1) => Boolean
        int jdbcType = introspectedColumn.getJdbcType();
        if(!tinyInt1isBit && Types.TINYINT == jdbcType) {
            if (introspectedColumn.getLength() == 1) {
                return new FullyQualifiedJavaType(Boolean.class.getName());
            } else {
                return new FullyQualifiedJavaType(Integer.class.getName());
            }
        }

        // 自定义类型映射
        JdbcTypeInformation jdbcTypeInformation = typeMap.get(jdbcType);
        if (jdbcTypeInformation == null) {

            switch (jdbcType) {
                case Types.DECIMAL:
                case Types.NUMERIC:
                    // 如果包含小数点则转换成float
                    if (introspectedColumn.getScale() > 0) {
                        result = new FullyQualifiedJavaType(Float.class.getName());
                    } else {
                        if (introspectedColumn.getLength() > 18 || forceBigDecimals) {
                            result = new FullyQualifiedJavaType(BigDecimal.class.getName());
                        } else if (introspectedColumn.getLength() > 9) {
                            result = new FullyQualifiedJavaType(Long.class.getName());
                        } else if (introspectedColumn.getLength() > 4) {
                            result = new FullyQualifiedJavaType(Integer.class.getName());
                        } else {
                            result = new FullyQualifiedJavaType(Short.class.getName());
                        }
                    }
                    break;
                default:
                    result = super.calculateJavaType(introspectedColumn);
                    break;
            }
        } else {
            result = super.calculateJavaType(introspectedColumn);
        }

        return result;
    }
}
